# -*- coding: utf-8 -*-
"""
Created on Wed Nov  4 16:17:24 2020

@author: Team317
"""

import torch
import torchvision.transforms as transforms
import torch.nn.functional as F
from age_net.resnet100 import KitModel
import PIL.Image as Image
import numpy as np
from face_modules.mtcnn import *
import cv2
import matplotlib.pyplot as plt
import pickle
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

# 计算人脸特征向量
def get_feature(img_path):
    
    # mtcnn检测人脸
    detector = MTCNN()
    original_img = cv2.imread(img_path)
    align_img = detector.align(Image.fromarray(original_img[:, :, ::-1]), crop_size=(256, 256))
    align_img = Image.fromarray(np.array(align_img))
    
    # 预处理
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    trans_img = test_transform(align_img)
    trans_img = trans_img.unsqueeze(0).cpu()
    
    # 用arcface计算人脸特征，优化：加载的过程可以放到函数外，避免多次加载
    arcface = KitModel('../Arcface_100.pth').to('cpu')
    arcface.eval()
    with torch.no_grad():
        after_handle = F.interpolate(trans_img[:, :, 19:237, 19:237], (112, 112), mode='bilinear', align_corners=True)
        embeds = arcface(after_handle)

    return  embeds.numpy()

# 计算两个人脸特征向量的距离
def distance(feature1, feature2):
    feature1 = np.array(feature1)
    feature2 = np.array(feature2)
    try:
        distance = np.linalg.norm(feature1[0] - feature2[0])
    except:
        # print(feature1,'\t',feature2)
        # print('type:', type(feature1), type(feature2))
        distance = 0
    return distance


    
# 计算同一个不同照片对的距离并记录
def same_face(f_list):
    n = len(f_list)
    D_list = []
    for i in range(n):
        for j in range(0,i):
            D = distance(f_list[i],f_list[j])
            if D != 0:
                D_list.append(D)
                
    return D_list
    
def diff_face(f_list1, f_list2):
    n1 = len(f_list1)
    n2 = len(f_list2)
    D_list = []
    for i in range(n1):
        for j in range(n2):
            D = distance(f_list1[i],f_list2[j])
            if D != 0:
                D_list.append(D)

    return D_list


if __name__ == '__main__':

    root = '../images/'
    folders = os.listdir(root);
    m = len(folders); n = 40
    faces_path = np.empty((m,n),dtype=object);
    
    i=0;    j=0;    data=dict()
    # 如果已保存特征向量，则不需要计算特征向量，置为False
    if False:
        # 获取每张图片的路径
        f = open('tempfile.txt', 'wb')
        for sub_folder in folders:
            folder = root + sub_folder
            faces = os.listdir(folder)
            temp = []
            for face in faces:
                face_path = folder + '/' + face
                # faces_path[i][j] = face_path
                try:
                    feature = get_feature(face_path)
                except:
                    print('invalid:',face_path)
                    feature = None
                temp.append(feature)
                j+=1
            data[str(i)] = temp
            i+=1
            j=0
            
        pickle.dump(data, f)
        f.close()
        
    
    f = open('tempfile.txt', 'rb')
    all_feature = pickle.load(f)
    
    # 计算同一个人不同照片的距离
    diff_D = []
    for i in range(1, len(all_feature)):
        for j in range(0, i):
            tempD = diff_face(all_feature[str(i)], all_feature[str(j)])
            diff_D.append(tempD)
    
    
    # 计算不同人之间的距离
    same_D = []
    for i in range(m):
        tempD = same_face(all_feature[str(i)])
        same_D.append(tempD)
        
    ## 绘制散点图 
    
    diff_data = [k for j in diff_D for k in j]
    diff_data1 = [0 for j in diff_data]
    p2= plt.scatter(diff_data, diff_data1, c='r', marker='1')
    
    same_data = [k for j in same_D for k in j]
    same_data1 = [0 for j in same_data]
    p1= plt.scatter(same_data, same_data1, c='b', marker='2')
    plt.legend((p1, p2), ('same person', 'difference person'))     #添加图例
    

    plt.show()
    
    


    
